#include "Common.fxh"
#include "ShadowFactor.fxh"

#define MAX_POINT_LIGHTS 128
#define MAX_SPOT_LIGHTS 128
#define MAX_SHD_SPOT_LIGHTS 16

#ifdef NO_SHADOWS
#else
Texture2D gCSMShadowMap;
#endif

cbuffer cbPerDirLight
{
	float3 gLightDirW;
	float4 gLightColorAndIntens;
#ifdef NO_SHADOWS
#else
	// Since in PS pixel's view space position is obtained, in order to draw shadows
	// it is necessary to transform it to light projection space, in following order:
	// view space -> world space -> light space -> light projection space
	// Entire transformations is in this matrix
	float4x4 gLightViewProjs[NUM_CASCADES];
	float4 gCascadeDepths;
#endif
};

struct PointLight
{
	float4 PosWAndRadius;
	float4 ColorAndIntens;
};

struct SpotLight
{
	float4 PosWAndRadius;
	float4 ColorAndIntens;
	float4 DirAndCosAngle;
};

struct ShdSpotLight
{
	float4 PosWAndRadius;
	float4 ColorAndIntens;
	float4 DirAndCosAngle;
	float4x4 ViewProj;
};

cbuffer cbPerLightSet
{
	PointLight gPointLights[MAX_POINT_LIGHTS];
	int gNumPointLights;

	SpotLight gSpotLights[MAX_SPOT_LIGHTS];
	int gNumSpotLights;

	ShdSpotLight gShdSpotLights[MAX_SHD_SPOT_LIGHTS];
	int gNumShdSpotLights;
};

Texture2D gSpotShadowMaps[MAX_SHD_SPOT_LIGHTS];

cbuffer cbPerFrame
{
	float4x4 gViewMatrix;
	float4x4 gProjMatrix;
	float3 gEyePos;
};

Texture2D gDiffuseMap;

struct VS_IN
{
	float3 posL : POSITION;
	float2 texCoord : TEXCOORD0;
	float size : TEXCOORD1;
	float opacity : TEXCOORD2;
};

struct VS_OUT_PER_PIXEL
{
	float4 posH : SV_POSITION;
	float3 posW : TEXCOORD0;
	float posV_z : TEXCOORD1;
	float2 texCoord : TEXCOORD2;
	float opacity : TEXCOORD3;
	float3 viewRight : TEXCOORD4;
	float3 viewUp : TEXCOORD5;
	float3 viewLook : TEXCOORD6;
};

struct VS_OUT_PER_VERTEX
{
	float4 posH : SV_POSITION;
	float2 texCoord : TEXCOORD0;
	float opacity : TEXCOORD1;
	float3 viewRight : TEXCOORD2;
	float3 viewUp : TEXCOORD3;
	float3 viewLook : TEXCOORD4;
	float3 hl2Light0 : TEXCOORD5;
	float3 hl2Light1 : TEXCOORD6;
	float3 hl2Light2 : TEXCOORD7;
};
 
//VS_OUT VS(VS_IN vIn)
//{
	//VS_OUT vOut;
//
	//float4 posV = mul(float4(vIn.posL, 1.0f), gWorldViewMatrix);
//
	//float2 vertexOffset = vIn.texCoord * float2(1, -1) + float2(-0.5, 0.5);
	//posV.xy += vIn.size * vertexOffset;
//
	//vOut.posH = mul(posV, gProjMatrix);
	//vOut.posW = mul(float4(vIn.posL, 1.0f), gWorldMatrix).xyz;
//
	//vOut.texCoord = vIn.texCoord;
	//vOut.opacity = vIn.opacity;
//
	//return vOut;
//}

VS_OUT_PER_VERTEX VS_PerVertex(VS_IN vIn)
{
	VS_OUT_PER_VERTEX vOut;

	float4 posW = float4(vIn.posL, 1.0f);

	float3 viewLook = normalize(posW.xyz - gEyePos);
	float3 viewRight = normalize(cross(float3(0,1,0), viewLook));
	float3 viewUp = cross(viewLook, viewRight);

	float2 vertexOffset = vIn.texCoord * float2(1, -1) + float2(-0.5, 0.5);
	posW.xyz += vIn.size * vertexOffset.x * viewRight;
	posW.xyz += vIn.size * vertexOffset.y * viewUp;

	float4 posV = mul(posW, gViewMatrix);
	vOut.posH = mul(posV, gProjMatrix);

	vOut.texCoord = vIn.texCoord;
	vOut.opacity = vIn.opacity;

	vOut.viewRight = viewRight;
	vOut.viewUp = viewUp;
	vOut.viewLook = viewLook;

	vOut.hl2Light0 = vOut.hl2Light1 = vOut.hl2Light2 = 0;

	float3x3 viewPosInv = float3x3(normalize(viewRight), normalize(viewUp), normalize(-viewLook));

	const float3 hl2Basis0 = float3(-1/sqrt(2), 1/sqrt(6), 1/sqrt(3));
	const float3 hl2Basis1 = float3( 1/sqrt(2), 1/sqrt(6), 1/sqrt(3));
	const float3 hl2Basis2 = float3(0, -sqrt(2/3), 1/sqrt(3));

	// dir light
	{
		float shadowFactor = 1.0f;

#ifdef NO_SHADOWS
#else
		if(posV.z > gCascadeDepths.z)
		{
			float4 lightSpacePos = mul(posW, gLightViewProjs[3]);
			shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.5, 0.5), lightSpacePos);
		}
		else if(posV.z > gCascadeDepths.y)
		{
			float4 lightSpacePos = mul(posW, gLightViewProjs[2]);
			shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.0, 0.5), lightSpacePos);
		}
		else if(posV.z > gCascadeDepths.x)
		{
			float4 lightSpacePos = mul(posW, gLightViewProjs[1]);
			shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.5, 0.0), lightSpacePos);
		}
		else
		{
			float4 lightSpacePos = mul(posW, gLightViewProjs[0]);
			shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.0, 0.0), lightSpacePos);
		}
#endif

		float3 lightColor = gLightColorAndIntens.rgb * gLightColorAndIntens.a;
		float3 lightDir = normalize(mul(-gLightDirW, viewPosInv));
		float3 weights = float3(saturate(dot(hl2Basis0, lightDir)),
								saturate(dot(hl2Basis1, lightDir)),
								saturate(dot(hl2Basis2, lightDir)));
		vOut.hl2Light0 += weights.x * (lightColor * shadowFactor + 0.0 * lightColor * (1 - shadowFactor));
		vOut.hl2Light1 += weights.y * (lightColor * shadowFactor + 0.0 * lightColor * (1 - shadowFactor));
		vOut.hl2Light2 += weights.z * (lightColor * shadowFactor + 0.0 * lightColor * (1 - shadowFactor));
	}

	for(int i = 0; i < gNumPointLights; i++)
	{
		const float3 lightPos = gPointLights[i].PosWAndRadius.xyz;
		const float lightRadius = gPointLights[i].PosWAndRadius.w;

		const float3 lightColor = gPointLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gPointLights[i].ColorAndIntens.a;

		float3 lightDir = lightPos - posW.xyz;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);
		float atten = saturate((lightRadius - dist) / lightRadius);

		float3 lighting = lightColor * lightIntensity * atten;
		float3 weights = float3(saturate(dot(hl2Basis0, lightDir)),
								saturate(dot(hl2Basis1, lightDir)),
								saturate(dot(hl2Basis2, lightDir)));
		vOut.hl2Light0 += weights.x * lighting;
		vOut.hl2Light1 += weights.y * lighting;
		vOut.hl2Light2 += weights.z * lighting;
	}

	for(int i = 0; i < gNumSpotLights; i++)
	{
		const float3 lightPos = gSpotLights[i].PosWAndRadius.xyz;
		const float lightRadius = gSpotLights[i].PosWAndRadius.w;

		const float3 lightColor = gSpotLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gSpotLights[i].ColorAndIntens.a;

		const float3 spotDir = gSpotLights[i].DirAndCosAngle.xyz;
		const float coneCosAngle = gSpotLights[i].DirAndCosAngle.w;

		float3 lightDir = lightPos - posW.xyz;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);

		float spotAngleCos = saturate(dot( -spotDir, lightDir ));

		if(spotAngleCos < coneCosAngle)
			continue;

		float atten = saturate((lightRadius - dist) / lightRadius);
		float angleAtten = saturate((spotAngleCos - coneCosAngle) / (1.0 - coneCosAngle));

		float3 lighting = lightColor * lightIntensity * atten * angleAtten;
		float3 weights = float3(saturate(dot(hl2Basis0, lightDir)),
								saturate(dot(hl2Basis1, lightDir)),
								saturate(dot(hl2Basis2, lightDir)));
		vOut.hl2Light0 += weights.x * lighting;
		vOut.hl2Light1 += weights.y * lighting;
		vOut.hl2Light2 += weights.z * lighting;
	}

	for(int i = 0; i < gNumShdSpotLights; i++)
	{
		const float3 lightPos = gShdSpotLights[i].PosWAndRadius.xyz;
		const float lightRadius = gShdSpotLights[i].PosWAndRadius.w;

		const float3 lightColor = gShdSpotLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gShdSpotLights[i].ColorAndIntens.a;

		const float3 spotDir = gShdSpotLights[i].DirAndCosAngle.xyz;
		const float coneCosAngle = gShdSpotLights[i].DirAndCosAngle.w;

		float3 lightDir = lightPos - posW.xyz;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);

		float spotAngleCos = saturate(dot( -spotDir, lightDir ));

		//if(spotAngleCos < coneCosAngle)
		//	continue;

		float atten = saturate((lightRadius - dist) / lightRadius);
		float angleAtten = saturate((spotAngleCos - coneCosAngle) / (1.0 - coneCosAngle));

		float4 lightSpacePos = mul(posW, gShdSpotLights[i].ViewProj);
		float shadowFactor = calculateShadowFactor(gSpotShadowMaps[i], lightSpacePos);

		float3 lighting = lightColor * lightIntensity * atten * angleAtten * shadowFactor;
		float3 weights = float3(saturate(dot(hl2Basis0, lightDir)),
								saturate(dot(hl2Basis1, lightDir)),
								saturate(dot(hl2Basis2, lightDir)));
		vOut.hl2Light0 += weights.x * lighting;
		vOut.hl2Light1 += weights.y * lighting;
		vOut.hl2Light2 += weights.z * lighting;
	}

	return vOut;
}

float4 PS_PerVertex(VS_OUT_PER_VERTEX pIn) : SV_Target
{
	float4 diffuseAndAlpha = gDiffuseMap.Sample(gTriLinearSam, pIn.texCoord);

	float3 normal = normalize(diffuseAndAlpha.xyz * 2 - 1);
	normal.x *= -1;

	const float3 hl2Basis0 = float3(-1/sqrt(2), 1/sqrt(6), 1/sqrt(3));
	const float3 hl2Basis1 = float3( 1/sqrt(2), 1/sqrt(6), 1/sqrt(3));
	const float3 hl2Basis2 = float3(0, -sqrt(2/3), 1/sqrt(3));

	float3 weights = float3(saturate(dot(normal, hl2Basis0)),
							saturate(dot(normal, hl2Basis1)),
							saturate(dot(normal, hl2Basis2)));
	float3 light = pIn.hl2Light0 * weights.x + pIn.hl2Light1 * weights.y + pIn.hl2Light0 * weights.z;

	return float4(light, diffuseAndAlpha.a * pIn.opacity);
}

VS_OUT_PER_PIXEL VS_PerPixel(VS_IN vIn)
{
	VS_OUT_PER_PIXEL vOut;

	float4 posW = float4(vIn.posL, 1.0f);

	float3 viewLook = normalize(posW.xyz - gEyePos);
	float3 viewRight = normalize(cross(float3(0,1,0), viewLook));
	float3 viewUp = cross(viewLook, viewRight);

	float2 vertexOffset = vIn.texCoord * float2(1, -1) + float2(-0.5, 0.5);
	posW.xyz += vIn.size * vertexOffset.x * viewRight;
	posW.xyz += vIn.size * vertexOffset.y * viewUp;

	float4 posV = mul(posW, gViewMatrix);
	vOut.posH = mul(posV, gProjMatrix);
	vOut.posW = posW.xyz;
	vOut.posV_z = posV.z;

	vOut.texCoord = vIn.texCoord;
	vOut.opacity = vIn.opacity;

	vOut.viewRight = viewRight;
	vOut.viewUp = viewUp;
	vOut.viewLook = viewLook;

	return vOut;
}

void addPointLights(inout float3 accum, in float3 posW, in float3 normalW)
{
	for(int i = 0; i < gNumPointLights; i++)
	{
		const float3 lightPos = gPointLights[i].PosWAndRadius.xyz;
		const float lightRadius = gPointLights[i].PosWAndRadius.w;

		const float3 lightColor = gPointLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gPointLights[i].ColorAndIntens.a;

		float3 lightDir = lightPos - posW;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);
		float atten = saturate((lightRadius - dist) / lightRadius);

		accum += saturate(dot(normalW, lightDir)) * lightColor * lightIntensity * atten;
	}
}

void addSpotLights(inout float3 accum, in float3 posW, in float3 normalW)
{
	for(int i = 0; i < gNumSpotLights; i++)
	{
		const float3 lightPos = gSpotLights[i].PosWAndRadius.xyz;
		const float lightRadius = gSpotLights[i].PosWAndRadius.w;

		const float3 lightColor = gSpotLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gSpotLights[i].ColorAndIntens.a;

		const float3 spotDir = gSpotLights[i].DirAndCosAngle.xyz;
		const float coneCosAngle = gSpotLights[i].DirAndCosAngle.w;

		float3 lightDir = lightPos - posW;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);

		float spotAngleCos = saturate(dot( -spotDir, lightDir ));

		if(spotAngleCos < coneCosAngle)
			continue;

		float atten = saturate((lightRadius - dist) / lightRadius);
		float angleAtten = saturate((spotAngleCos - coneCosAngle) / (1.0 - coneCosAngle));

		accum += saturate(dot(normalW, lightDir)) * lightColor * lightIntensity * atten * angleAtten;
	}
}

void addShdSpotLights(inout float3 accum, in float3 posW, in float3 normalW)
{
	for(int i = 0; i < gNumShdSpotLights; i++)
	{
		const float3 lightPos = gShdSpotLights[i].PosWAndRadius.xyz;
		const float lightRadius = gShdSpotLights[i].PosWAndRadius.w;

		const float3 lightColor = gShdSpotLights[i].ColorAndIntens.rgb;
		const float lightIntensity = gShdSpotLights[i].ColorAndIntens.a;

		const float3 spotDir = gShdSpotLights[i].DirAndCosAngle.xyz;
		const float coneCosAngle = gShdSpotLights[i].DirAndCosAngle.w;

		float3 lightDir = lightPos - posW;
		float dist = length(lightDir);
		lightDir = normalize(lightDir);

		float spotAngleCos = saturate(dot( -spotDir, lightDir ));

		//if(spotAngleCos < coneCosAngle)
		//	continue;

		float atten = saturate((lightRadius - dist) / lightRadius);
		float angleAtten = saturate((spotAngleCos - coneCosAngle) / (1.0 - coneCosAngle));

		float4 lightSpacePos = mul(float4(posW, 1.0f), gShdSpotLights[i].ViewProj);
		float shadowFactor = calculateShadowFactor(gSpotShadowMaps[i], lightSpacePos);

		accum += saturate(dot(normalW, lightDir)) * lightColor * lightIntensity * atten * angleAtten * shadowFactor;
	}
}

float4 PS_PerPixel(VS_OUT_PER_PIXEL pIn) : SV_Target
{
	float4 diffuseAndAlpha = gDiffuseMap.Sample(gTriLinearSam, pIn.texCoord);

	float3 normal = normalize(diffuseAndAlpha.xyz * 2 - 1);
	normal.x *= -1;

	//float3 viewLook = normalize(pIn.posW - gEyePos);
	//float3 viewRight = normalize(cross(float3(0,1,0), viewLook));
	//float3 viewUp = cross(viewLook, viewRight);
	float3x3 viewPosInv = float3x3(normalize(pIn.viewRight), normalize(pIn.viewUp), normalize(-pIn.viewLook));
	
	normal = normalize(mul(normal, viewPosInv));
	//normal = mul(float4(normal, 0), gInvViewMatrix).xyz; // transform normal from view to world space
	//normal = normalize(gEyePos - pIn.posW);

	float3 lightAccum = 0;

	// dir light

	float shadowFactor = 1.0f;

#ifdef NO_SHADOWS
#else
	if(pIn.posV_z > gCascadeDepths.z)
	{
		float4 lightSpacePos = mul(float4(pIn.posW, 1.0f), gLightViewProjs[3]);
		shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.5, 0.5), lightSpacePos);
	}
	else if(pIn.posV_z > gCascadeDepths.y)
	{
		float4 lightSpacePos = mul(float4(pIn.posW, 1.0f), gLightViewProjs[2]);
		shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.0, 0.5), lightSpacePos);
	}
	else if(pIn.posV_z > gCascadeDepths.x)
	{
		float4 lightSpacePos = mul(float4(pIn.posW, 1.0f), gLightViewProjs[1]);
		shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.5, 0.0), lightSpacePos);
	}
	else
	{
		float4 lightSpacePos = mul(float4(pIn.posW, 1.0f), gLightViewProjs[0]);
		shadowFactor = calculateShadowFactorCSM(gCSMShadowMap, float2(0.0, 0.0), lightSpacePos);
	}
#endif

	float3 dirLighting = gLightColorAndIntens.rgb * gLightColorAndIntens.a * saturate(dot(normal, -gLightDirW));
	lightAccum += dirLighting * shadowFactor;// + 0.5 * dirLighting * (1 - shadowFactor);

	addPointLights(lightAccum, pIn.posW, normal);
	addSpotLights(lightAccum, pIn.posW, normal);
	addShdSpotLights(lightAccum, pIn.posW, normal);

	return float4(lightAccum, diffuseAndAlpha.a * pIn.opacity);
}

technique11 Layout
{
	pass P0
	{
	#ifdef PER_VERTEX
		SetVertexShader( CompileShader( vs_5_0, VS_PerVertex() ) );
		SetPixelShader( CompileShader( ps_5_0, PS_PerVertex() ) );
	#else
		SetVertexShader( CompileShader( vs_5_0, VS_PerPixel() ) );
		SetPixelShader( CompileShader( ps_5_0, PS_PerPixel() ) );
	#endif

		SetGeometryShader( NULL );

		SetDepthStencilState(DepthWriteOff, 0);
		SetBlendState( Composite, float4( 0.0f, 0.0f, 0.0f, 0.0f ), 0xFFFFFFFF );
	}
}
